package com.db.concurrency;

import java.math.BigInteger;
import java.util.concurrent.Callable;
import java.util.concurrent.CancellationException;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Future;
import java.util.concurrent.FutureTask;

public class Memoizer<A, V> implements Computable<A, V> {
  private final Computable<A, V>                computer;
  private final ConcurrentHashMap<A, Future<V>> cache = new ConcurrentHashMap<>();

  public Memoizer(Computable<A, V> computer) {
    this.computer = computer;
  }

  @Override
  public V compute(final A arg) throws InterruptedException {
    while (true) {
      Future<V> f = cache.get(arg);
      if (f == null) {
        Callable<V> eval = new Callable<V>() {
          public V call() throws InterruptedException {
            return computer.compute(arg);
          }
        };
        FutureTask<V> ft = new FutureTask<V>(eval);
        f = cache.putIfAbsent(arg, ft);
        if (f == null) {
          f = ft;
          ft.run(); // call to computer.compute happens here
        }
      }
      try {
        return f.get();
      } catch (CancellationException e) {
        cache.remove(arg);
      } catch (ExecutionException e) {
        cache.remove(arg);
        throw new RuntimeException(e);
      }
    }
  }
}

class ExpensiveComputation implements Computable<String, BigInteger> {

  @Override
  public BigInteger compute(String arg) throws InterruptedException {
    return new BigInteger(arg);
  }

}